import argparse
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy import interpolate
from scipy.stats import linregress
from pathlib import Path
import sys

# 设置全局字体为支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义高对比度颜色调色板
color_list = [
    '#FF0000',  # 鲜红色
    '#0000FF',  # 鲜蓝色
    '#800080',  # 紫红色
    '#FF8000',  # 橙色
    '#8000FF',  # 紫色
    '#FF0080',  # 玫红色
    '#80FF00',  # 黄绿色
    '#0080FF',  # 天蓝色
    '#FF8080',  # 浅红色
    '#800000',  # 深红色
    '#008000',  # 深绿色
    '#000080',  # 深蓝色
    '#808000',  # 橄榄色
    '#008080',  # 深青色
    '#FFA500',  # 金橙色
    '#32CD32',  # 石灰绿
    '#FF1493',  # 深粉色
    '#1E90FF',  # 道奇蓝
]

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='多文件数据读取和直线拟合工具')

    parser.add_argument('input_files', nargs='+',
                       help='输入的CSV文件列表')
    parser.add_argument('-x', '--x_column', required=True,
                       help='X轴数据字段名')
    parser.add_argument('-y', '--y_columns', nargs='+', required=True,
                       help='Y轴数据字段名列表')
    parser.add_argument('-f', '--flags', nargs='*',
                       help='文件标识符列表，用于区分不同文件的数据。如果不指定，默认使用文件名')
    parser.add_argument('-t', '--title',
                       help='图表标题。如果不指定，将自动生成标题')
    
    # 添加坐标轴标签参数
    parser.add_argument('--xlabel',
                       help='X轴标签名称。如果不指定，使用X轴字段名')
    parser.add_argument('--ylabel',
                       help='Y轴标签名称。如果不指定，使用默认的"value"')
    
    parser.add_argument('--output', '-o', default='output_plot.png',
                       help='输出图片文件名 (默认: output_plot.png)')
    parser.add_argument('--figsize', nargs=2, type=float, default=[14, 8],
                       help='图片尺寸 (宽度 高度)')

    return parser.parse_args()

def validate_arguments(args):
    """验证参数的有效性"""
    if args.flags is not None and len(args.flags) != len(args.input_files):
        raise ValueError(f"flags参数数量({len(args.flags)})必须与输入文件数量({len(args.input_files)})相同")

    return True

def load_and_merge_data(file_paths, x_column, y_columns, flags=None):
    """加载并合并多个CSV文件的数据，为Y轴字段添加文件标识"""
    all_data = []
    final_y_columns = []  # 存储最终的Y轴字段名

    for idx, file_path in enumerate(file_paths):
        if not Path(file_path).exists():
            print(f"警告: 文件 {file_path} 不存在，跳过")
            continue

        try:
            df = pd.read_csv(file_path)
        except Exception as e:
            print(f"警告: 无法读取文件 {file_path}: {e}，跳过")
            continue

        # 检查必需的列是否存在
        required_columns = [x_column] + y_columns
        missing_columns = [col for col in required_columns if col not in df.columns]
        if missing_columns:
            print(f"警告: 文件 {file_path} 缺少列: {missing_columns}，跳过")
            continue

        # 确定文件标识
        file_flag = flags[idx] if flags is not None else Path(file_path).stem

        # 选择需要的数据
        selected_data = df[[x_column]].copy()

        # 为每个Y轴字段添加文件标识前缀
        renamed_cols = []
        for y_col in y_columns:
            new_y_col = f"{file_flag}.{y_col}"
            selected_data[new_y_col] = df[y_col]
            renamed_cols.append(new_y_col)
            if new_y_col not in final_y_columns:
                final_y_columns.append(new_y_col)

        selected_data['source_file'] = file_flag
        all_data.append(selected_data)

        print(f"✓ 成功加载文件: {file_path}")
        print(f"  - 文件标识: {file_flag}")
        print(f"  - 数据行数: {len(df)}")
        print(f"  - Y字段重命名: {renamed_cols}")

    # 直接按行拼接
    merged_data = pd.concat(all_data, ignore_index=True)

    # 按x轴排序
    merged_data = merged_data.sort_values(by=x_column).reset_index(drop=True)

    print(f"✓ 合并后总数据行数: {len(merged_data)}")
    print(f"✓ 最终Y轴字段: {final_y_columns}")
    print(f"✓ 数据按 {x_column} 字段排序完成")

    return merged_data, final_y_columns

def perform_linear_regression(x_data, y_data):
    """执行线性回归分析"""
    # 移除NaN值
    mask = ~(np.isnan(x_data) | np.isnan(y_data))
    x_clean = x_data[mask]
    y_clean = y_data[mask]

    if len(x_clean) < 2:
        raise ValueError("有效数据点不足，无法进行拟合")

    slope, intercept, r_value, p_value, std_err = linregress(x_clean, y_clean)

    return {
        'slope': slope,
        'intercept': intercept,
        'r_value': r_value,
        'r_squared': r_value**2,
        'p_value': p_value,
        'std_err': std_err,
        'x_clean': x_clean,
        'y_clean': y_clean
    }

def calculate_r_squared(y_actual, y_predicted):
    """计算R²值"""
    ss_res = np.sum((y_actual - y_predicted) ** 2)
    ss_tot = np.sum((y_actual - np.mean(y_actual)) ** 2)
    if ss_tot == 0:
        return 1.0 if ss_res == 0 else 0.0
    return 1 - (ss_res / ss_tot)

def print_regression_results(results_dict, x_column, final_y_columns):
    """打印回归分析结果"""
    print("\n" + "=" * 60)
    print("线性拟合结果")
    print("=" * 60)

    for y_col in final_y_columns:
        if y_col in results_dict:
            result = results_dict[y_col]
            print(f"\n{y_col} vs {x_column} 拟合结果:")
            print(f"拟合直线方程: Y = {result['slope']:.4f} * X + {result['intercept']:.4f}")
            print(f"斜率: {result['slope']:.4f}")
            print(f"截距: {result['intercept']:.4f}")
            print(f"相关系数 R: {result['r_value']:.4f}")
            print(f"决定系数 R²: {result['r_squared']:.4f}")
            print(f"P值: {result['p_value']:.6f}")
            print(f"标准误差: {result['std_err']:.4f}")
            print(f"有效数据点数: {len(result['x_clean'])}")

def create_comprehensive_plot(data, results_dict, x_column, final_y_columns, args):
    """创建综合图表"""
    fig, ax = plt.subplots(figsize=args.figsize)

    # 使用自定义颜色列表，如果不够则循环使用
    num_colors_needed = len(final_y_columns)
    colors = [color_list[i % len(color_list)] for i in range(num_colors_needed)]

    for idx, y_col in enumerate(final_y_columns):
        if y_col not in results_dict:
            continue

        result = results_dict[y_col]
        color = colors[idx]

        # 获取该y列的非NaN数据
        y_data = data[y_col].dropna()
        if len(y_data) == 0:
            continue

        x_data = data.loc[y_data.index, x_column]

        # 按x值排序
        sort_idx = np.argsort(x_data)
        x_sorted = x_data.iloc[sort_idx]
        y_sorted = y_data.iloc[sort_idx]

        # 1. 绘制原始数据点
        ax.scatter(x_sorted, y_sorted, color=color, marker='o', s=100,
                  label=f'{y_col} (原始数据)', zorder=5, alpha=0.8)

        # 2. 绘制相邻点虚线连接
        ax.plot(x_sorted, y_sorted, color=color, linestyle='--',
               linewidth=2, alpha=0.6, label=f'{y_col} (插值数据)')

        # 3. 绘制拟合直线
        x_range = np.linspace(x_sorted.min(), x_sorted.max(), 100)
        y_fitted = result['slope'] * x_range + result['intercept']

        ax.plot(x_range, y_fitted, color=color, linestyle='-',
               linewidth=3, alpha=0.9,
               label=f'{y_col} 拟合标准差:{result["std_err"]:.2f}\n拟合直线:Y = {result["slope"]:.2f}X + {result["intercept"]:.2f}')

    # 设置轴标签 - 使用新添加的参数
    xlabel = args.xlabel if args.xlabel else x_column
    ylabel = args.ylabel if args.ylabel else 'value'
    
    ax.set_xlabel(xlabel, fontsize=14)
    ax.set_ylabel(ylabel, fontsize=14)

    # 设置标题
    if args.title:
        title = args.title
    else:
        title = f'{", ".join(final_y_columns)} 随 {x_column} 变化曲线'
    ax.set_title(title, fontsize=16)

    # 设置图例位置为右上角
    ax.legend(fontsize=9, loc='upper right', frameon=True,
             fancybox=True, shadow=True, framealpha=0.9)

    # 添加网格
    ax.grid(True, linestyle='--', alpha=0.7)

    # 应用tight_layout
    plt.tight_layout()

    # 保存图片
    plt.savefig(args.output, dpi=300, bbox_inches='tight')
    plt.show()

    print(f"\n✓ 图表已保存到: {args.output}")

def print_data_statistics(data, x_column, final_y_columns):
    """打印数据统计信息"""
    print("\n" + "=" * 60)
    print("数据范围统计")
    print("=" * 60)

    print(f"{x_column} 值范围: {data[x_column].min():.2f} 到 {data[x_column].max():.2f}")

    for y_col in final_y_columns:
        if y_col in data.columns:
            y_data = data[y_col].dropna()
            if len(y_data) > 0:
                print(f"{y_col} 值范围: {y_data.min():.2f} 到 {y_data.max():.2f}")

    print(f"\n数据来源文件: {', '.join(data['source_file'].unique())}")

def create_sample_data():
    """创建示例数据文件用于测试"""
    # 创建示例数据文件1
    np.random.seed(42)
    x1 = np.linspace(0, 10, 20)
    y1_1 = 2 * x1 + 1 + np.random.normal(0, 0.5, len(x1))
    y1_2 = -0.5 * x1 + 8 + np.random.normal(0, 0.3, len(x1))

    df1 = pd.DataFrame({
        'time': x1,
        'temperature': y1_1,
        'pressure': y1_2
    })
    df1.to_csv('sample_data1.csv', index=False)

    # 创建示例数据文件2
    x2 = np.linspace(5, 15, 15)
    y2_1 = 1.8 * x2 + 2 + np.random.normal(0, 0.4, len(x2))
    y2_2 = -0.3 * x2 + 7 + np.random.normal(0, 0.2, len(x2))

    df2 = pd.DataFrame({
        'time': x2,
        'temperature': y2_1,
        'pressure': y2_2
    })
    df2.to_csv('sample_data2.csv', index=False)

    print("✓ 示例数据文件已创建: sample_data1.csv, sample_data2.csv")

def main():
    """主函数"""
    try:
        args = parse_arguments()
        validate_arguments(args)

        print("=== 多文件数据拟合工具 ===")
        print(f"输入文件: {args.input_files}")
        print(f"X轴字段: {args.x_column}")
        print(f"Y轴字段: {args.y_columns}")
        print(f"文件标识: {args.flags if args.flags else '使用文件名'}")
        print(f"图表标题: {args.title if args.title else '自动生成'}")
        print(f"X轴标签: {args.xlabel if args.xlabel else f'使用字段名({args.x_column})'}")
        print(f"Y轴标签: {args.ylabel if args.ylabel else '数值'}")

        # 加载和合并数据
        merged_data, final_y_columns = load_and_merge_data(
            args.input_files, args.x_column, args.y_columns, args.flags)

        # 对每个最终的y字段进行线性回归
        results_dict = {}

        for y_col in final_y_columns:
            try:
                # 获取非NaN的数据
                y_data = merged_data[y_col].dropna()
                if len(y_data) == 0:
                    print(f"警告: 字段 {y_col} 没有有效数据，跳过")
                    continue

                x_data = merged_data.loc[y_data.index, args.x_column]

                result = perform_linear_regression(x_data.values, y_data.values)
                results_dict[y_col] = result

            except Exception as e:
                print(f"警告: 无法处理字段 {y_col}: {e}")
                continue

        if not results_dict:
            raise ValueError("没有成功处理任何Y轴字段")

        # 打印回归结果
        print_regression_results(results_dict, args.x_column, final_y_columns)

        # 打印数据统计
        print_data_statistics(merged_data, args.x_column, final_y_columns)

        # 创建图表
        create_comprehensive_plot(merged_data, results_dict, args.x_column, final_y_columns, args)

        print("\n" + "=" * 60)
        print("拟合质量评估")
        print("=" * 60)

        for y_col, result in results_dict.items():
            y_pred = result['slope'] * result['x_clean'] + result['intercept']
            r2 = calculate_r_squared(result['y_clean'], y_pred)
            print(f"{y_col} 拟合优度 R²: {r2:.4f}")

        print("\n=== 处理完成 ===")

    except Exception as e:
        print(f"错误: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
